from sklearn.svm import SVC
from sklearn import datasets
from sklearn.model_selection import train_test_split as split
import numpy as np
from matplotlib import pyplot as plt
import matplotlib.animation as animation


# ==================================
#         编写实现TSVM算法的函数
# ==================================
def TSVM(Dl, Du, Cl, Cu, gifname=None):
    '''
    # 该函数实现TSVM算法，其中支持向量机算法基于sklearn中的SVC类的线性支持向量机实现
    #
    # 输入参数：
    #     Dl：有标记样本数据,格式为元祖形式：(Xl,yl),其中yl取值为{1,-1}
    #     Du：未标记样本数据，仅Xu
    #     Cl,Cu:初始折中参数，Cu<<Cl
    #     gifname:若要将TSVM算法过程保存为gif动画，则传入文件名，默认不保存动画
    # 输出：
    #     clf：基于sklearn的支持向量机SVM分类器
    '''

    Xl, yl = Dl
    Xu = Du
    X_mix = np.r_[Xl, Xu]
    clf = SVC(C=Cl, kernel='linear').fit(Xl, yl)  # 基于有标记样本的初始SVM分类器
    yu = clf.predict(Xu)  # 向未标记样本指派伪标记
    # acts用于后续绘制动画所用，其中储存了相应的事件动作的关键参数，
    # 比如：重新拟合后的权重w和偏置b、重新分派伪标记后的伪标记yu等
    acts = [{'w': clf.coef_.reshape(-1), 'b': clf.intercept_, 'text': '初始模型(仅有标记样本)'}]
    acts.append({'assign': yu.copy(), 'text': '分派伪标记'})
    while Cu < Cl:
        # 样本权重，传入clf.fit()函数可实现对于不同样本不同的权重值
        sample_weight = [1.0] * len(yl) + [Cu / Cl] * len(yu)
        clf.fit(X_mix, np.r_[yl, yu], sample_weight=sample_weight)
        acts.append({'w': clf.coef_.reshape(-1), 'b': clf.intercept_, 'text': '调整Cu为%.3E后重新拟合' % Cu})
        while True:
            f = clf.decision_function(Xu)  # 计算f(x)=wx+b的结果
            xi = np.fmax(1 - f * yu, 0)  # 计算ξ值,基于y(wx+b)≥1-ξ,ξ≥0
            y1_index = np.where(yu == 1)[0]  # 伪标记为+1的索引号
            y0_index = np.where(yu == -1)[0]  # 伪标记为-1的索引号
            max1 = max(xi[yu == 1])  # 伪标记为+1的样本中的最大ξ值
            max0 = max(xi[yu == -1])  # 伪标记为-1的样本中的最大ξ值
            # 只需分别考虑伪正负样本中最大ξ值的两个样本即可，
            # 因为若这两个最大ξ值的样本不满足条件(ξi>0,ξj>0,ξi+ξj>2)，
            # 那么其他样本对也必然无法满足了。
            if (max1 > 0) & (max0 > 0) & (max1 + max0 > 2):
                print('交换伪标记：ξ_+1=%.3f,ξ_-1=%.3f' % (max1, max0))
                i = y1_index[np.argmax(xi[yu == 1])]  # 伪标记为+1的样本中的最大ξ值对应的样本索引号
                j = y0_index[np.argmax(xi[yu == -1])]  # 伪标记为-1的样本中的最大ξ值对应的样本索引号
                yu[i] *= -1
                yu[j] *= -1
                acts.append({'exchanging': [i, j], 'text': '交换伪标记中...'})
                acts.append({'assign': yu.copy(), 'text': '完成交换'})
                clf.fit(X_mix, np.r_[yl, yu], sample_weight=sample_weight)
                acts.append({'w': clf.coef_.reshape(-1), 'b': clf.intercept_, 'text': '交换标记后重新拟合'})
            else:
                break
        Cu = min(2 * Cu, Cl)
    acts.append({'text': 'TSVM算法执行完毕!'})

    if gifname != None:
        # 设置绘图中显示中文
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False

        fig = plt.figure()

        # 绘制有标记样本，用颜色区分正负样本
        plt.scatter(Xl[yl == 1, 0], Xl[yl == 1, 1], s=40, c='r', marker='+', edgecolors='r')
        plt.scatter(Xl[yl == -1, 0], Xl[yl == -1, 1], s=40, c='g', marker='_', edgecolors='g')
        # 绘制有标记样本，全部显示为黑色小点
        plt.scatter(Xu[:, 0], Xu[:, 1], s=10, c='k')
        # 设置坐标轴上下限
        x1min, x1max = min(X_mix[:, 0]), max(X_mix[:, 0])
        x2min, x2max = min(X_mix[:, 1]), max(X_mix[:, 1])
        plt.xlim([x1min - (x1max - x1min) * 0.2, x1max + (x1max - x1min) * 0.2])
        plt.ylim([x2min - (x2max - x2min) * 0.2, x2max + (x2max - x2min) * 0.2])

        # 分类器决策线
        decision_line, = plt.plot([], [], c='k')
        # 无标记样本之+1样本指派结果，颜色与有标记样本对应颜色相同
        unlabel_points1, = plt.plot([], [], marker='.', linestyle='', c='r')
        # 无标记样本之-1样本指派结果，颜色与有标记样本对应颜色相同
        unlabel_points0, = plt.plot([], [], marker='.', linestyle='', c='g')
        # 绘制交换标记操作的牵引线
        exchange, = plt.plot([0, 0], [0, 0], linestyle='', c='m', linewidth=2)  # 初始，任意位置，设置为不显示
        # 显示当前操作状态的文字说明
        state = plt.text((x1min + x1max) / 2, x2max + (x2max - x2min) * 0.1, '', fontsize=12)

        def update(num):
            # 更新动画帧的函数，num为帧数
            act = acts[num]
            if 'w' in act:
                decision_line.set_data([x1min, x1max],
                                       [-(x1min * act['w'][0] + act['b']) / act['w'][1],
                                        -(x1max * act['w'][0] + act['b']) / act['w'][1]])
                exchange.set_linestyle('')
                state.set_text(act['text'])
            elif 'exchanging' in act:
                i, j = act['exchanging']
                exchange.set_linestyle('--')
                exchange.set_data([Xu[i][0], Xu[j][0]], [Xu[i][1], Xu[j][1]])
                state.set_text(act['text'])
            elif 'assign' in act:
                yu = act['assign']
                unlabel_points1.set_data(Xu[yu == 1, 0], Xu[yu == 1, 1])
                unlabel_points0.set_data(Xu[yu == -1, 0], Xu[yu == -1, 1])
                state.set_text(act['text'])
            else:
                state.set_text(act['text'])

            return [decision_line, unlabel_points1, unlabel_points0, exchange, state]

        # 动画生成函数animation.FuncAnimation(),其中参数intervel为每帧的持续时间，单位为ms
        ani = animation.FuncAnimation(fig, update, frames=range(len(acts)), interval=1000)
        # 关于保存gif动画，有些电脑可能直接可以正常运行，
        # 有些电脑则会报错，这是因为电脑系统中缺少了某些组件，
        # 可以根据提示以及参考网络上一些案例进行安装，
        # 比如，我参考了下列网址中的方法成功地解决了问题：
        # https://blog.csdn.net/weixin_41957054/article/details/107280246
        # https://blog.csdn.net/qq_21905401/article/details/103023074
        ani.save(str(gifname) + '.gif', writer='imagemagick')
        plt.show()
    return clf


# ==========================================
#     生成简单二维数据集，
#     明显线性可分离的两类数据，
#     试验并观察TSVM算法过程
#     计算过程通过gif动画的形式进行演示
# ==========================================
# X1和X2为生成的两类数据
X1 = np.random.random([50, 2])
X2 = np.random.random([50, 2]) + [2, 0]
X1 = X1[np.argsort(X1[:, 1])]
X2 = X2[np.argsort(X2[:, 1])]
# 在X1和X2中各取五个样本组成有标记样本
Xl = np.r_[X1[:5], X2[-5:]]
yl = np.array([1] * 5 + [-1] * 5)
# 其余的样本作为无标记样本
Xu = np.r_[X1[5:], X2[:-5]]
TSVM((Xl, yl), Xu, 1, 0.0001, 'demo1')

# ==========================================
#     同样生成简单二维数据集，
#     但是两类数据有重合，无法明显线性划分
# ==========================================
# X1和X2为生成的两类数据
X1 = np.random.random([50, 2])
X2 = np.random.random([50, 2]) + [1, 0]
X1 = X1[np.argsort(X1[:, 1])]
X2 = X2[np.argsort(X2[:, 1])]
# 在X1和X2中各取五个样本组成有标记样本
Xl = np.r_[X1[:5], X2[-5:]]
yl = np.array([1] * 5 + [-1] * 5)
# 其余的样本作为无标记样本
Xu = np.r_[X1[5:], X2[:-5]]
TSVM((Xl, yl), Xu, 1, 0.0001, 'demo2')

# ==========================================
#    在莺尾花数据集上进行试验
# ==========================================
print('-------在莺尾花数据集上进行试验-------')
iris = datasets.load_iris()
X = iris['data']
y = (iris['target'] == 0) * 2 - 1  # 将第1类设为y=+1,第2、3类设为y=-1

# 划分数据集:X_test:Xu:Xl样本比例为：3:6:1
X_train, X_test, y_train, y_test = split(X, y, test_size=0.3, random_state=12)
Xu, Xl, yu, yl = split(X_train, y_train, test_size=1 / 7, random_state=12)

# 单独有标记样本进行训练和预测
Cl = 1
clf0 = SVC(C=Cl, kernel='linear').fit(Xl, yl)
y_pre0 = clf0.predict(X_test)
acc0 = (y_pre0 == y_test).mean()
print('仅用有标记样本进行训练的模型的预测精度为：', acc0)

# 利用无标记样本的半监督模型
Cl, Cu = 1, 0.0001
clf1 = TSVM((Xl, yl), Xu, Cl, Cu)
y_pre1 = clf1.predict(X_test)
acc1 = (y_pre1 == y_test).mean()
print('利用无标记样本的半监督模型的预测精度为：', acc1)
print('预测精度提高了%.2f%%' % ((acc1 - acc0) / acc0 * 100))

# ==========================================
#    在手写数据集上进行试验
# ==========================================
print('-------在手写数据集上进行试验-------')
digits = datasets.load_digits()
X = digits['data']
y = (digits['target'] < 5) * 2 - 1  # 将0~4设为y=+1,第5~9设为y=-1

# 划分数据集:X_test:Xu:Xl样本比例为：3:6:1
X_train, X_test, y_train, y_test = split(X, y, test_size=0.3, random_state=12)
Xu, Xl, yu, yl = split(X_train, y_train, test_size=1 / 7, random_state=12)

# 单独有标记样本进行训练和预测
Cl = 1
clf0 = SVC(C=Cl, kernel='linear').fit(Xl, yl)
y_pre0 = clf0.predict(X_test)
acc0 = (y_pre0 == y_test).mean()
print('仅用有标记样本进行训练的模型的预测精度为：', acc0)

# 利用无标记样本的半监督模型
Cl, Cu = 1, 0.0001
clf1 = TSVM((Xl, yl), Xu, Cl, Cu)
y_pre1 = clf1.predict(X_test)
acc1 = (y_pre1 == y_test).mean()
print('利用无标记样本的半监督模型的预测精度为：', acc1)
print('预测精度提高了%.2f%%' % ((acc1 - acc0) / acc0 * 100))
